import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据预处理：转换为张量 + 归一化
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转为张量（形状：[1,28,28]）
    transforms.Normalize((0.1307,), (0.3081,))  # 归一化（MNIST数据集的均值和标准差）
])

# 加载训练集和测试集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 数据加载器（批量加载数据，加速训练）
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积：输入通道1（灰度图），输出通道16，卷积核3×3
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()  #ReLU 激活函数，增加非线性特征提取能力
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 池化后尺寸：14×14

        # 第二层卷积：输入通道16，输出通道32，卷积核3×3
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 池化后尺寸：7×7

        # 全连接层：将卷积特征展平后分类（32×7×7 → 10类）
        self.fc = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        # 第一层卷积流程
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        # 第二层卷积流程
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)

        # 展平特征 → 全连接层分类
        x = x.view(-1, 32 * 7 * 7)
        x = self.fc(x)
        return x


model = CNN()  # 初始化模型

criterion = nn.CrossEntropyLoss()  # 交叉熵损失（适用于分类任务）
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器，学习率0.001

epochs = 5
for epoch in range(epochs):
    model.train()  # 切换为训练模式
    train_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()  # 清零梯度
        output = model(data)  # 模型前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()  # 反向传播（计算梯度）
        optimizer.step()  # 更新参数
        train_loss += loss.item()

        # 每100个batch打印一次训练进度
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch + 1}/{epochs}, Batch {batch_idx}, Loss: {loss.item():.4f}')
    # 打印当前epoch的平均训练损失
    print(f'Epoch {epoch + 1} Average Train Loss: {train_loss / len(train_loader):.4f}')

    model.eval()  # 切换为评估模式（关闭 dropout、BatchNorm等）
    test_loss = 0
    correct = 0
    with torch.no_grad():  # 禁用梯度计算（加速推理）
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)  # 预测类别（取概率最大的类别）
            correct += pred.eq(target.view_as(pred)).sum().item()  # 统计正确数

    # 计算测试集损失和准确率
    test_loss = test_loss / len(test_loader)
    test_accuracy = correct / len(test_loader.dataset)
    print(f'Test Loss: {test_loss:.4f}')
    print(f'Test Accuracy: {test_accuracy:.4f}')  # 通常可达到98%以上